import os

import pandas as pd
import torch
import torch.nn as nn
from pathlib import Path
from torchmetrics import Accuracy, ConfusionMatrix
from openpyxl.reader.excel import load_workbook


def print_progress_bar(iteration, total, prefix='', suffix='', decimals=1, length=100, fill='█', print_end="\r"):
    """
    调用在Python终端中打印自定义进度条的功能
    iteration - 当前的迭代（Int）
    total - 迭代的总数（Int）
    prefix - 前缀字符串（Str）
    suffix - 后缀字符串（Str）
    decimals - 正数的小数位数（Int）
    length - 进度条的长度（Int）
    fill - 进度条填充字符（Str）
    print_end - 行尾字符（Str）
    """
    percent = ("{0:." + str(decimals) + "f}").format(100 * (iteration / float(total)))
    filled_length = int(length * iteration // total)
    bar = fill * filled_length + '-' * (length - filled_length)
    print(f'\r{prefix} |{bar}| {percent}% {suffix}', end=print_end)


def print_log(mdic, epoch, all_batch, batch_id, loss, acc):
    """
       该函数可以用于打印模型训练过程中的日志
       参数说明：
       Epoch:训练值
       All_batch:总批次
       Batch_num:训练到的批次数
       loss_value:损失函数值
       acc_value:评价指标
       step_time:一个批次的执行时间，单位为毫秒
    """
    p = int(batch_id * 100 / all_batch)
    process = [">"] * p + [" "] * (100 - p)
    mdic["epoch"].append(epoch)
    mdic["batch_id"].append(batch_id)
    mdic["loss"].append(loss)
    mdic["acc"].append(acc)
    print("\r Epoch: {}, batch: {}, process:{},loss is: {:.4f}, acc is: {:.4f}".format(
        epoch,
        str(batch_id) + "/" + str(
            all_batch),
        ''.join(
            process),
        loss,
        acc),
        end="")


def save_log_local(mdic, model_name, is_training=True):
    # 将字典保存为xlsx
    df = pd.DataFrame(mdic)

    # 根据是训练还是测试数据设置文件名

    if is_training:

        filename = "训练日志数据.xlsx"

        dir_path = "日志数据\\训练"  # 注意这里的路径分隔符

    else:

        filename = "测试日志数据.xlsx"

        dir_path = "日志数据\\测试"  # 注意这里的路径分隔符

    # 创建目录，如果它不存在

    os.makedirs(dir_path, exist_ok=True)  # `exist_ok=True` 防止目录已存在时引发异常

    # 构建完整的文件路径

    file_path = os.path.join(dir_path, filename)

    if Path(file_path).exists():
        # 加载现有的工作簿
        book = load_workbook(file_path)

        # 如果只有一个工作表则删除整个文件
        if len(book.sheetnames) == 1 and book.sheetnames[0] == model_name:
            os.remove(file_path)
            # 如果文件不存在，则直接写入
            with pd.ExcelWriter(file_path, engine='openpyxl') as writer:
                df.to_excel(writer, sheet_name=model_name, index=False)
        else:
            # 如果工作表存在，则删除
            if model_name in book.sheetnames:
                del book[model_name]
            book.save(file_path)
            # 重新打开临时文件以写入新的数据
            with pd.ExcelWriter(file_path, engine='openpyxl', mode='a') as writer:
                # 写入新的 DataFrame 到指定的工作表
                df.to_excel(writer, sheet_name=model_name, index=False)
    else:
        # 如果文件不存在，则直接写入
        with pd.ExcelWriter(file_path, engine='openpyxl') as writer:
            df.to_excel(writer, sheet_name=model_name, index=False)


class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        # 设置当前模型最多支持多少个批次
        self.max_batch = 0
        # 模型名
        self.model_name = ""

    def my_train(self, epochs, data_loader, data_loader_eval):
        """
        训练模型
        :param epochs: 训练轮数
        :param data_loader: 数据加载器
        :return:
        """
        # 初始化self.max_batch
        if self.max_batch == 0:
            self.max_batch = len(data_loader)
        # 设置优化器和损失函数
        optimizer = torch.optim.AdamW(self.parameters(), lr=0.0001)
        criterion = torch.nn.CrossEntropyLoss()
        # 定义acc
        # 初始化准确度评估器
        acc_metric = Accuracy(task='multiclass', num_classes=5).to(self.device)

        # 假设 outputs 是模型的预测输出，labels 是真实的标签

        # 创建一个空字典 用于记录日志
        mdic = {"epoch": [], "batch_id": [], "loss": [], "acc": []}
        # 创建一个列表记录每一次的epoch_acc
        epoch_accs = []

        for epoch in range(epochs):
            self.train()
            for i, (images, labels) in enumerate(data_loader):
                # 重置评估器以便下次使用
                acc_metric.reset()

                # 将images与labels 按照 self.max_batch 进行分割
                num_samples = images.shape[0]

                num_splits = (num_samples + self.max_batch - 1) // self.max_batch  # Calculate the number of splits
                outputs = torch.tensor([]).to(self.device)
                total_loss = 0.0

                for j in range(num_splits):
                    start_idx = j * self.max_batch

                    end_idx = min((j + 1) * self.max_batch, num_samples)

                    # Get the current split.

                    current_images = images[start_idx:end_idx]

                    current_labels = labels[start_idx:end_idx]
                    # 前向传播
                    current_outputs = self(current_images.to(self.device))
                    outputs = torch.cat((outputs, current_outputs), dim=0)
                    #
                    current_loss = criterion(current_outputs, current_labels.to(self.device)).to(self.device)
                    total_loss += current_loss.item()
                    # 反向传播和优化
                    current_loss.backward()

                # 执行更新
                optimizer.step()
                optimizer.zero_grad()
                # 在每个批次或整个数据集上更新评估器
                acc_metric.update(outputs, labels.to(self.device))
                # 计算并获取当前准确度
                accuracy = acc_metric.compute()
                # 打印损失
                print_log(mdic, epoch, len(data_loader), i, total_loss / num_splits, accuracy.item())
            epoch_acc = self.my_eval(data_loader_eval)
            epoch_accs.append(epoch_acc)
            # 判断一下epoch_accs 最近五个 epoch_acc是否还在增长 如果没有就停止训练。
            if len(epoch_accs) > 5:
                # 如果epoch_accs的最大值不在最后2个数中
                if max(epoch_accs[-3:]) != max(epoch_accs):
                    break

        # 将数据保存为xsl并保存到本地
        save_log_local(mdic, self.model_name)
        # 保存模型
        print(" ")
        # 利用pytorch保存现有模型所有内容
        torch.save(self.state_dict(), "模型参数\\" + self.model_name + ".pth")
        print(" 训练结束,保存模型参数成功")

    def my_eval(self, data_loader):
        self.eval()
        acc_metric = Accuracy(task='multiclass', num_classes=5).to(self.device)
        all_acc = 0
        for i, (images, labels) in enumerate(data_loader):
            outputs = self(images)
            # 在每个批次或整个数据集上更新评估器
            acc_metric.update(outputs, labels.to(self.device))
            # 计算并获取当前准确度
            accuracy = acc_metric.compute()
            all_acc += accuracy.item()
        return all_acc / (len(data_loader))

    def my_test(self, data_loader):
        self.eval()
        mdic = {"epoch": [], "batch_id": [], "loss": [], "acc": []}  # 创建一个空字典 用于记录日志
        acc_metric = Accuracy(task='multiclass', num_classes=5).to(self.device)
        confusion_matrix = ConfusionMatrix(task='multiclass', num_classes=5).to(self.device)
        all_acc = 0
        all_confusion_matrix = None  # 初始化混淆矩阵
        for i, (images, labels) in enumerate(data_loader):
            outputs = self(images)
            loss = 0  # 如果有需要，可以在这里计算损失
            # 在每个批次或整个数据集上更新评估器
            acc_metric.update(outputs, labels.to(self.device))
            # 更新混淆矩阵
            confusion_matrix.update(outputs, labels.to(self.device))
            # 计算并获取当前准确度
            accuracy = acc_metric.compute()
            all_acc += accuracy.item()
            print_log(mdic, -1, len(data_loader), i, loss, accuracy.item())
        print(" 测试集准确率:", all_acc / (len(data_loader)))
        # 输出每个类别的分类情况
        print(" 混淆矩阵:")
        # 获取混淆矩阵
        all_confusion_matrix = confusion_matrix.compute()
        print(all_confusion_matrix)
        # 将混淆矩阵转换为DataFrame
        df_confusion_matrix = pd.DataFrame(all_confusion_matrix.cpu().numpy(), columns=[f"Class_{i}" for i in range(5)],
                                           index=[f"Class_{i}" for i in range(5)])
        # 判断是否存在混淆矩阵文件夹
        if not os.path.exists("日志数据\\混淆矩阵"):
            os.makedirs("日志数据\\混淆矩阵")
        # 判断是否存在"日志数据\\测试\\混淆矩阵\\{self.model_name}_confusion_matrix.csv" 存在则删除
        if os.path.exists(f"日志数据\\混淆矩阵\\{self.model_name}_confusion_matrix.csv"):
            os.remove(f"日志数据\\混淆矩阵\\{self.model_name}_confusion_matrix.csv")
        # 保存混淆矩阵到CSV文件
        df_confusion_matrix.to_csv(f"日志数据\\混淆矩阵\\{self.model_name}_confusion_matrix.csv")
        # 保存日志
        save_log_local(mdic, self.model_name, is_training=False)
